import SimpleITK as sitk 
import numpy as np 
import pydicom
import pydicom_seg

template = pydicom_seg.template.from_dcmqi_metainfo('metainfo.json')

writer = pydicom_seg.MultiClassWriter(
    template=template,
    inplane_cropping=False,  # Crop image slices to the minimum bounding box on
                            # x and y axes
    skip_empty_slices=False,  # Don't encode slices with only zeros
    skip_missing_segment=False,  # If a segment definition is missing in the
                                 # template, then raise an error instead of
                                 # skipping it.
)

reader = sitk.ImageSeriesReader()
dcm_files = reader.GetGDCMSeriesFileNames('./out_image/', '1.2.826.0.1.3680043.2.1125.20230610.1231229')
reader.SetFileNames(dcm_files)
image = reader.Execute()
image_data = sitk.GetArrayFromImage(image)

print(image_data.shape)

segmentation_data = np.zeros_like(image_data).astype(np.uint8)

segmentation_data[10:200] = 1

segmentation = sitk.GetImageFromArray(segmentation_data)
segmentation.CopyInformation(image)

source_images = [
    pydicom.dcmread(x, stop_before_pixels=False)
    for x in dcm_files
]
dcm = writer.write(segmentation, source_images)

dcm.save_as('segmentation.dcm')
